package com.coder.algorithm.offer_2;

import com.coder.algorithm.struct.BinaryTreeNode;

import java.util.Stack;

/**
 * 二叉搜索树的第 k 大节点
 *
 * @author yuhushuan
 * @date 2020/6/30 09:22
 */
public class _54_1_KthNodeInBST {

    private static BinaryTreeNode getKthNode(BinaryTreeNode root, int k) {
        if (root == null) {
            return null;
        }

        BinaryTreeNode node = root;
        Stack<BinaryTreeNode> stack = new Stack<>();
        while (node != null || !stack.isEmpty()) {
            while (node != null) {
                stack.push(node);
                node = node.left;
            }
            node = stack.pop();
            if (k-- == 1) {
                return node;
            }
            node = node.right;
        }
        return null;
    }

    private static BinaryTreeNode getKthNode_Recursively(BinaryTreeNode root, int k) {
        if (root == null) {
            return null;
        }
        return getKthNodeCore_Recursively(root, new int[]{k});
    }

    private static BinaryTreeNode getKthNodeCore_Recursively(BinaryTreeNode node, int[] k) {
        if (node == null) {
            return null;
        }
        BinaryTreeNode kth = null;
        if (node.left != null) {
            kth = getKthNodeCore_Recursively(node.left, k);
        }
        if (kth != null) {
            return kth;
        }
        if (k[0]-- == 1) {
            kth = node;
        }
        if (kth == null && node.right != null) {
            kth = getKthNodeCore_Recursively(node.right, k);
        }
        return kth;
    }


    public static void main(String[] args) {
        int len = 7;
        BinaryTreeNode[] nodes = new BinaryTreeNode[7];
        // {2,3,4,5,6,7,8}
        for (int i = 0; i < len; ++i) {
            nodes[i] = new BinaryTreeNode(i + 2);
        }
        nodes[3].left = nodes[1];
        nodes[3].right = nodes[5];

        nodes[1].left = nodes[0];
        nodes[1].right = nodes[2];

        nodes[5].left = nodes[4];
        nodes[5].right = nodes[6];

        // 功能测试
        print(nodes[3], 3);

        // 边界值测试
        print(nodes[3], 0);
        print(nodes[3], 1);
        print(nodes[3], 7);
        print(nodes[3], 8);

        // 特殊输入
        print(null, 1);
    }

    private static void print(BinaryTreeNode root, int k) {
        System.out.println("二叉搜索树：");
        BinaryTreeNode.drawBinaryTree(root);
        BinaryTreeNode node = getKthNode(root, k);
        System.out.println("第 " + k + " 大节点为：" + (node == null ? "null" : node.value));
        node = getKthNode_Recursively(root, k);
        System.out.println("第 " + k + " 大节点为(递归)：" + (node == null ? "null" : node.value));
    }
}
